#!/usr/bin/python
# -*- coding: utf-8 -*-
# pylint: disable=C0103
# pylint: disable=E1101

import sys
import time
import numpy as np
import tensorflow as tf
import cv2

from utils import label_map_util
from utils import visualization_utils_color as vis_util
# 检测模型的路径
PATH_TO_CKPT = './model/frozen_inference_graph_face.pb'

# 标签映射文件的路径
PATH_TO_LABELS = './protos/face_label_map.pbtxt'

NUM_CLASSES = 2
# 加载标签映射文件
label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
# 将标签映射转换为类别字典
categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES,
                                                            use_display_name=True)
# 创建类别索引
category_index = label_map_util.create_category_index(categories)

class TensoflowFaceDector(object):
    def __init__(self, PATH_TO_CKPT):
        """Tensorflow detector
        """
        """Tensorflow 人脸检测器初始化
                """
        # 加载检测模型
        self.detection_graph = tf.Graph()
        with self.detection_graph.as_default():
            od_graph_def = tf.compat.v1.GraphDef()
            with tf.io.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
                serialized_graph = fid.read()
                od_graph_def.ParseFromString(serialized_graph)
                tf.import_graph_def(od_graph_def, name='')
        # 配置 TensorFlow 会话
        with self.detection_graph.as_default():
            config = tf.compat.v1.ConfigProto()
            config.gpu_options.allow_growth = True
            self.sess = tf.compat.v1.Session(graph=self.detection_graph, config=config)
            self.windowNotSet = True

    def run(self, image):
        """image: bgr image
        return (boxes, scores, classes, num_detections)
        """
        """图像中检测人脸
                Args:
                    image: BGR图像
                Returns:
                    (boxes, scores, classes, num_detections)
                """
        # 将图像从BGR转换为RGB
        image_np = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        # the array based representation of the image will be used later in order to prepare the
        # result image with boxes and labels on it.
        # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
        # 将图像扩展维度以符合模型输入要求
        image_np_expanded = np.expand_dims(image_np, axis=0)
        image_tensor = self.detection_graph.get_tensor_by_name('image_tensor:0')
        # Each box represents a part of the image where a particular object was detected.
        # 获取模型输出节点
        boxes = self.detection_graph.get_tensor_by_name('detection_boxes:0')
        # Each score represent how level of confidence for each of the objects.
        # Score is shown on the result image, together with the class label.
        scores = self.detection_graph.get_tensor_by_name('detection_scores:0')
        classes = self.detection_graph.get_tensor_by_name('detection_classes:0')
        num_detections = self.detection_graph.get_tensor_by_name('num_detections:0')
        # Actual detection.
        # 进行目标检测
        start_time = time.time()
        (boxes, scores, classes, num_detections) = self.sess.run(
            [boxes, scores, classes, num_detections],
            feed_dict={image_tensor: image_np_expanded})
        elapsed_time = time.time() - start_time
        print('inference time cost: {}'.format(elapsed_time))

        return (boxes, scores, classes, num_detections)


if __name__ == "__main__":
    import sys

    if len(sys.argv) != 2:
        print("usage:%s (cameraID | filename) Detect faces\
 in the video example:%s 0" % (sys.argv[0], sys.argv[0]))
        exit(1)

    try:
        camID = int(sys.argv[1])
    except:
        camID = sys.argv[1]
    # 初始化人脸检测器
    tDetector = TensoflowFaceDector(PATH_TO_CKPT)
    # 打开摄像头或视频文件
    cap = cv2.VideoCapture(camID)
    windowNotSet = True
    while True:
        ret, image = cap.read()
        if ret == 0:
            break

        [h, w] = image.shape[:2]
        print(h, w)
        # 翻转图像
        image = cv2.flip(image, 1)
        # 运行人脸检测器
        (boxes, scores, classes, num_detections) = tDetector.run(image)
        # 可视化检测结果
        vis_util.visualize_boxes_and_labels_on_image_array(
            image,
            np.squeeze(boxes),
            np.squeeze(classes).astype(np.int32),
            np.squeeze(scores),
            category_index,
            use_normalized_coordinates=True,
            line_thickness=4)

        if windowNotSet is True:
            cv2.namedWindow("tensorflow based (%d, %d)" % (w, h), cv2.WINDOW_NORMAL)
            windowNotSet = False

        cv2.imshow("tensorflow based (%d, %d)" % (w, h), image)
        k = cv2.waitKey(1) & 0xff
        if k == ord('q') or k == 27:
            break

    cap.release()

"""
这段代码实现了使用 TensorFlow 模型进行人脸检测的功能。
首先初始化了 TensorFlow 人脸检测器，然后从摄像头或视频文件中读取帧图像，对图像进行预处理后，
调用 TensorFlow 人脸检测器进行推断，最后将检测结果可视化并展示出来。
"""